#!/usr/bin/env python3
# -*- coding:UTF8 -*-

import cv2 as cv
import numpy as np
import os

"""
从一张带有车牌中的图片中扣出车牌
"""

"""
形态学初步处理
"""
def preprocess(img):
     # 高斯平滑
    gaussian = cv.GaussianBlur(img, (3, 3), 0, 0, cv.BORDER_DEFAULT)
    median = cv.medianBlur(gaussian, 5)
    
    gray = cv.cvtColor(median, cv.COLOR_BGR2GRAY)
    # Soble 算子 X方向梯度
    sobel = cv.Sobel(gray, cv.CV_8UC1, 1, 0, cv.BORDER_DEFAULT)
    # 二值化
    ret, binary = cv.threshold(sobel, 170, 255, cv.THRESH_BINARY)
    # 膨胀腐蚀的核函数
    element1 = cv.getStructuringElement(cv.MORPH_RECT, (9, 3))
    elementcv = cv.getStructuringElement(cv.MORPH_RECT, (9, 5))
    # 膨胀一次，让轮廓突出
    dilation = cv.dilate(binary, elementcv, iterations=1)
    # 腐蚀一次，去掉细节
    erosion = cv.erode(dilation, element1, iterations=1)
    # 再次膨胀，让轮廓明显一些
    dilationcv = cv.dilate(erosion, elementcv, iterations=3)
    return dilationcv


# 保存疑似车牌的矩形区域
def findPlateNumberRegion(gray):
    region = []
    # 查找轮廓
    contours, hierarchy = cv.findContours(gray, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[-2:]

    # 筛选面积小的
    for i in range(len(contours)):
        cnt = contours[i]
        # 计算该轮廓的面积
        area = cv.contourArea(cnt)

        # 面积小的都筛选掉
        if (area < 200):
            continue

        # 找到最小的矩形，该矩形可能有方向
        rect = cv.minAreaRect(cnt)

        # box是四个点的坐标
        box = cv.boxPoints(rect)
        box = np.int0(box)

        # 计算高和宽
        height = abs(box[0][1] - box[2][1])
        width = abs(box[0][0] - box[2][0])
        # 车牌正常情况下长高比在7-5之间
        ratio = float(width) / float(height)
        if (ratio > 5 or ratio < 2):
            continue
        region.append(box)
    return region


def detect(img):
    thr = preprocess(img)
    cv.imshow("thr", thr)
    reg = findPlateNumberRegion(thr)
    for box in reg:
        ys = [box[0, 1], box[1, 1], box[2, 1], box[3, 1]]
        xs = [box[0, 0], box[1, 0], box[2, 0], box[3, 0]]
        ys_sorted_index = np.argsort(ys)
        xs_sorted_index = np.argsort(xs)

        x1 = box[xs_sorted_index[0], 0]
        x2 = box[xs_sorted_index[3], 0]

        y1 = box[ys_sorted_index[0], 1]
        y2 = box[ys_sorted_index[3], 1]
        img_orgcv = img.copy()
        img_plate = img_orgcv[y1:y2, x1:x2]
        return img_plate


def main():
    imagePath = './download'
    for img_path in os.listdir(imagePath):
        img_path = os.path.join(imagePath, img_path)
        img = cv.imread(img_path)
        img_plate = detect(img)
        if isinstance(img_plate, np.ndarray):
            cv.imshow('number plate', img_plate)
        cv.waitKey(0)


if __name__ == '__main__':
    main()
